rm(list = ls())
options(stringsAsFactors = FALSE)
seed_to_use <- 614
set.seed(seed_to_use)
library(data.table)
library(parallel)
library(ggplot2)
library(xtable)
#library(rstanarm) 
load("no_trade_msm.rdata")
## poisson coefs ----
coefs_list_pois <- (lapply(1:length(boot_msm_vals), function(i){
  # cat(i, "\n")
  if(!(is.null(boot_msm_vals[[i]]))){
  test_error <- (grepl(x = boot_msm_vals[[i]], pattern = "Error"))
  if(!test_error){
  as.data.table(t(boot_msm_vals[[i]]$quasi_poisson_weights))
  }
  }
}))
coefs_dt_pois <- rbindlist(coefs_list_pois)
y_00 <- coefs_dt_pois$`(Intercept)`

y_11 <- coefs_dt_pois$treatment + coefs_dt_pois$lag_1_treatment
mean(y_11)
y11_lower <- mean(y_11)  - 1.96 * sd(y_11)
y11_upper <-  mean(y_11)  + 1.96 * sd(y_11)

y_10 <- coefs_dt_pois$treatment
mean(y_10)
y_10_lower <- mean(y_10)  - 1.96 * sd(y_10)
y_10_upper <-  mean(y_10)  + 1.96 * sd(y_10)

y_01 <- coefs_dt_pois$lag_1_treatment
mean(y_01)
y_01_lower <- mean(y_01)  - 1.96 * sd(y_01)
y_01_upper <-  mean(y_01)  + 1.96 * sd(y_01)

## zi poisson coefs ----
coefs_list_zi_pois <- (lapply(1:length(boot_msm_vals), function(i){
 # cat(i, "\n")
  if(!(is.null(boot_msm_vals[[i]]))){
  test_error <- (grepl(x = boot_msm_vals[[i]], pattern = "Error"))
  if(!test_error){
  as.data.table(t(boot_msm_vals[[i]]$zero_infl_pois))
  }
  }
}))
coefs_list_zi_pois <- rbindlist(coefs_list_zi_pois)
y_11_zi <- coefs_list_zi_pois$count_treatment + coefs_list_zi_pois$count_lag_1_treatment
mean(y_11_zi)
y_11_zi_lower <- mean(y_11_zi)  - 1.96 * sd(y_11_zi)
y_11_zi_upper <-  mean(y_11_zi)  + 1.96 * sd(y_11_zi)

y_10_zi <- coefs_list_zi_pois$count_treatment
mean(y_10_zi)
quantile(y_10_zi, probs = c(.025, .975))
y_10_zi_lower <- mean(y_10_zi)  - 1.96 * sd(y_10_zi)
y_10_zi_upper <-  mean(y_10_zi)  + 1.96 * sd(y_10_zi)

y_01_zi <- coefs_list_zi_pois$count_lag_1_treatment
mean(y_01_zi)
quantile(y_01_zi, probs = c(.025, .975))
y_01_zi_lower <- mean(y_01_zi)  - 1.96 * sd(y_01_zi)
y_01_zi_upper <-  mean(y_01_zi)  + 1.96 * sd(y_01_zi)

y_11_zi_zero <- coefs_list_zi_pois$zero_treatment + coefs_list_zi_pois$zero_lag_1_treatment
mean(y_11_zi_zero)
y_11_zi_zero_lower <- mean(y_11_zi_zero)  - 1.96 * sd(y_11_zi_zero)
y_11_zi_zero_upper <-  mean(y_11_zi_zero)  + 1.96 * sd(y_11_zi_zero)

y_10_zi_zero <- coefs_list_zi_pois$zero_treatment
mean(y_10_zi_zero)
quantile(y_10_zi_zero, probs = c(.025, .975))
y_10_zi_zero_lower <- mean(y_10_zi_zero)  - 1.96 * sd(y_10_zi_zero)
y_10_zi_zero_upper <-  mean(y_10_zi_zero)  + 1.96 * sd(y_10_zi_zero)

y_01_zi_zero <- coefs_list_zi_pois$zero_lag_1_treatment
mean(y_01_zi_zero)
quantile(y_01_zi_zero, probs = c(.025, .975))
y_01_zi_zero_lower <- mean(y_01_zi_zero)  - 1.96 * sd(y_01_zi_zero)
y_01_zi_zero_upper <-  mean(y_01_zi_zero)  + 1.96 * sd(y_01_zi_zero)

poisson_dt <- data.table(
  model = c("Poisson Model"), 
  treatment = c("Liberalization at Time T and T-1", "Liberalization at Time T", "Liberalization at Time T-1"), 
  ATE = c(mean(y_11), mean(y_10), mean(y_01)), 
  lower_ci = c(quantile(y_11, .025), quantile(y_10, .025), quantile(y_01, .025)),
  upper_ci = c(quantile(y_11, .975), quantile(y_10, .975), quantile(y_01, .975)),
  lower_90 = c(quantile(y_11, .05), quantile(y_10, .05), quantile(y_01, .05)),
  upper_90 = c(quantile(y_11, .95), quantile(y_10, .95), quantile(y_01, .95)))

poisson_zi_dt <- data.table(
  model = c("Zero Inflated Poisson Model"), 
  treatment =c("Liberalization at Time T and T-1", "Liberalization at Time T", "Liberalization at Time T-1"), 
  ATE = c(mean(y_11_zi), mean(y_10_zi), mean(y_01_zi)), 
  lower_ci = c(quantile(y_11_zi, .025), quantile(y_10_zi, .025), quantile(y_01_zi, .025)),
  upper_ci = c(quantile(y_11_zi, .975), quantile(y_10_zi, .975), quantile(y_01_zi, .975)),
  lower_90 = c(quantile(y_11_zi, .05), quantile(y_10_zi, .05), quantile(y_01_zi, .05)),
  upper_90 = c(quantile(y_11_zi, .95), quantile(y_10_zi, .95), quantile(y_01_zi, .95)))

poisson_zi_zero_dt <- data.table(
  model = c("Zero Inflated Component Model"), 
  treatment =c("Liberalization at Time T and T-1", "Liberalization at Time T", "Liberalization at Time T-1"), 
  ATE = c(mean(y_11_zi_zero), mean(y_10_zi_zero), mean(y_01_zi_zero)), 
  lower_ci = c(quantile(y_11_zi_zero, .025), quantile(y_10_zi_zero, .025), quantile(y_01_zi_zero, .025)),
  upper_ci = c(quantile(y_11_zi_zero, .975), quantile(y_10_zi_zero, .975), quantile(y_01_zi_zero, .975)),
  lower_90 = c(quantile(y_11_zi_zero, .05), quantile(y_10_zi_zero, .05), quantile(y_01_zi_zero, .05)),
  upper_90 = c(quantile(y_11_zi_zero, .95), quantile(y_10_zi_zero, .95), quantile(y_01_zi_zero, .95)))

poisson_results <- rbindlist(list(poisson_dt, poisson_zi_dt))
poisson_results[, treatment := factor(treatment, 
  levels = c("Liberalization at Time T and T-1",  "Liberalization at Time T",  "Liberalization at Time T-1"))]
save(poisson_results, file = "poisson_msm_results_no_trade.rdata")

save(poisson_zi_zero_dt, file = "zero_inflated_results_no_trade.rdata")

p <- ggplot(data = poisson_results, aes(x = ATE, y = treatment)) + 
  geom_vline(aes(xintercept = 0), color = "grey80") + 
  geom_errorbarh(aes(xmax = upper_ci, xmin = lower_ci, y = treatment), height = .05, lwd = .5) + 
  geom_errorbarh(aes(xmax = upper_90, xmin = lower_90, y = treatment), height = 0, lwd = 1.5, color = "grey50") +
  geom_point(size = 4) + labs(title = "Treatment Binarized at Mean") +
  facet_wrap(~model) + theme(legend.position = "none", 
    panel.background = element_rect(fill = "white"),
    panel.grid.major = element_blank(), 
    panel.grid.minor = element_blank(), 
    legend.key = element_blank(), 
    legend.title = element_blank(),
    axis.line = element_blank(),
    axis.text.x = element_text(size=8),
    axis.text.y = element_text(size=10),
    plot.title = element_text(size=10),
    axis.title.y = element_blank(),
    text = element_text(family = "Palatino")) 
ggsave(plot = p, filename = "mean_binarized.pdf", device = "pdf")

